import os
import sys
import time
import pyexiv2
import hashlib

from stat import ST_SIZE

class JpegAnalysis(object):
    
    def __init__(self, path = None):
        ''' Constructor '''
        if path == None:
            self.path = os.environ['HOME'] + '/Pictures'
        else:
            self.path = path
        self.fileNames = []
        self.makeDictionary = {}
        self.modelDictionary = {}
        self.softwareDictionary = {}
        self.dateTimeDictionary = {}
        self.userCommentDictionary = {}
        self.fileSizeDictionary = {}
        self.hashList = []
        
    def analyzeFiles(self):
        ''' Analyzes the files in the self.path directory '''
        self.getFileNames()
        self.fillMetadataDictionaries()
    
    def getFileNames(self):
        ''' Get enumerate a directory, get file names only '''
        print ' [*] Enumerating:', self.path
        for entry in os.listdir(self.path):
            if os.path.isfile(self.path +'/'+ entry):
                self.fileNames.append(entry)
        print ' [*] Found %d file(s) to examine' % len(self.fileNames)
    
    def checkHash(self,imageName):
        ''' Calculate an Md5 hash for image '''
        targetFile = open(self.path +'/'+ imageName, 'r')
        md5Calculator = hashlib.md5()
        md5Calculator.update(targetFile.read())
        currentHash = md5Calculator.hexdigest()
        for hash in self.hashList:
            if hash == currentHash:
                return True
        self.hashList.append(currentHash)
        return False
    
    def getMakeField(self, tag):
        '''Get Make Field'''
        if self.makeDictionary.has_key(tag):
            self.makeDictionary[tag] += 1
        else:
            self.makeDictionary[tag] = 1

    def getModelField(self,tag):
        '''Get Model Field'''
        if self.modelDictionary.has_key(tag):
            self.modelDictionary[tag] += 1
        else:
            self.modelDictionary[tag] = 1

    def getSoftwareField(self,tag):
        '''Get Software Field'''
        if self.softwareDictionary.has_key(tag):
            self.softwareDictionary[tag] += 1
        else:
            self.softwareDictionary[tag] = 1

    def getDateTimeField(self,tag):
        '''Gets Date and Time of Creation'''
        if self.dateTimeDictionary.has_key(tag):
            self.dateTimeDictionary[tag] += 1
        else:
            self.dateTimeDictionary[tag] = 1

    def getUserCommentField(self,tag):
        '''Gets User Comment Field'''
        if self.userCommentDictionary.has_key(tag):
            self.userCommentDictionary[tag] += 1
        else:
            self.userCommentDictionary[tag] = 1
          
    def getFileSize(self, path):
        '''Gets File Size'''
        filestats = os.stat(path)
        fileSize = filestats[ST_SIZE]
        fileSize = (fileSize/8)/1024  # Convert to Kb
        if self.fileSizeDictionary.has_key(fileSize):
            self.fileSizeDictionary[fileSize] += 1
        else:
            self.fileSizeDictionary[fileSize] = 1

    def fillMetadataDictionaries(self):
        '''Fills the Statistic Dictionaries'''
        init = time.time()
        sys.stdout.write(' [*] Reading metadata: ')
        completed = 0
        for imageName in self.fileNames:
            progress = str(int(float(completed)/len(self.fileNames)*100))+'%'
            sys.stdout.write(progress)
            sys.stdout.flush()
            if self.checkHash(imageName):
                sys.stdout.write('\b \b' * len(progress))
                continue
            try:
                image = pyexiv2.ImageMetadata(self.path +'/'+ imageName)
                image.read()
                self.readExivData(image)
                self.getFileSize(self.path +'/'+ imageName)
                completed += 1
            except IOError:
                sys.stdout.write('\b \b' * len(progress))
                continue
            time.sleep(0.01)
            sys.stdout.write('\b \b' * len(progress))
            sys.stdout.flush()
        elapsed = time.time()-init
        sys.stdout.write('done!\n [*] Completed file analysis in %d second(s)\n' % elapsed)
    
    def readExivData(self, image):
        ''' Looks for exif data in the images '''
        for exifKey in image.exif_keys:
            if exifKey == 'Exif.Image.Make':
                tag = image['Exif.Image.Make']
                self.getMakeField(tag)
            if exifKey == 'Exif.Image.Model':
                tag = image['Exif.Image.Model']
                self.getModelField(tag)
            if exifKey == 'Exif.Image.Software':
                tag = image['Exif.Image.Software']
                self.getSoftwareField(tag)
            if exifKey == 'Exif.Image.Model':
                tag = image['Exif.Image.DateTime']
                self.getDateTimeField(tag)
            if exifKey == 'Exif.Photo.UserComment':
                tag = image['Exif.Photo.UserComment']
                self.getUserCommentField(tag)
                
    def getDictionary(self, name):
        if name == 'make':
            return self.makeDictionary, "Camera Make", "Manufacturor", "# Found"
        elif name == 'model':
            return self.modelDictionary, "Camera Model", "Model Name", "# Found"
        elif name == 'software':
            return self.softwareDictionary, "Photo Software", "Software Name", "# Found"
        elif name == 'date' or name == 'time':
            return self.dateTimeDictionary, "Date / Time", "", ""
        elif name == 'comment':
            return self.userCommentDictionary, "User Comment", "", ""
        elif name == 'size':
            return self.fileSizeDictionary, "Image File Size", "Kilobytes", "Number of Files"
        else:
            return None

if __name__ == '__main__':
    jp = JpegAnalysis()
    jp.analyzeFiles()
    